{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from segmentation_models_pytorch import Unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.tensor([\n",
    "    [0.2, 0.3, 0.4, 0.5],\n",
    "    [0.6, 0.6, 0.6, 0.6]\n",
    "])\n",
    "b = torch.tensor([\n",
    "    [0.1, 0.4, 0.3, 0.6],\n",
    "    [0.5, 0.7, 0.8, 0.9]\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 2, 4])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.stack([a], dim=0).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2000, 0.4000, 0.4000, 0.6000],\n",
       "        [0.6000, 0.7000, 0.8000, 0.9000]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stack = torch.stack([a, b], dim=0)\n",
    "stack.max(dim=0)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand((4, 3, 256, 256))\n",
    "y = (torch.rand((4, 3, 256, 256)) > 0.5).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 4, 3, 256, 256])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x_ = x.unsqueeze(0)\n",
    "x_ = x_.repeat(3, 1, 1, 1, 1)\n",
    "x_.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bce_2d(outputs, targets):\n",
    "    return F.binary_cross_entropy_with_logits(outputs.view(-1), targets.view(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.7336)\n",
      "tensor(0.7336)\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    print(bce_2d(x, y))\n",
    "    print(F.binary_cross_entropy_with_logits(x, y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load result: None\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Unet(\n",
       "  (encoder): EfficientNetEncoder(\n",
       "    (_conv_stem): Conv2dStaticSamePadding(\n",
       "      3, 32, kernel_size=(3, 3), stride=(2, 2), bias=False\n",
       "      (static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)\n",
       "    )\n",
       "    (_bn0): BatchNorm2d(32, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "    (_blocks): ModuleList(\n",
       "      (0): MBConvBlock(\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          32, 32, kernel_size=(3, 3), stride=[1, 1], groups=32, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(32, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          32, 8, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          8, 32, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(16, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (1): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(96, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          96, 96, kernel_size=(3, 3), stride=[2, 2], groups=96, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(96, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          96, 4, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          4, 96, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(24, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (2): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          144, 144, kernel_size=(3, 3), stride=(1, 1), groups=144, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          144, 6, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          6, 144, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          144, 24, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(24, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (3): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          144, 144, kernel_size=(5, 5), stride=[2, 2], groups=144, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 2, 1, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(144, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          144, 6, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          6, 144, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          144, 40, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(40, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (4): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          240, 240, kernel_size=(5, 5), stride=(1, 1), groups=240, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          240, 10, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          10, 240, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          240, 40, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(40, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (5): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          40, 240, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          240, 240, kernel_size=(3, 3), stride=[2, 2], groups=240, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(0, 1, 0, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(240, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          240, 10, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          10, 240, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          240, 80, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(80, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (6): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          480, 480, kernel_size=(3, 3), stride=(1, 1), groups=480, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          480, 20, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          20, 480, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          480, 80, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(80, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (7): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          480, 480, kernel_size=(3, 3), stride=(1, 1), groups=480, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          480, 20, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          20, 480, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          480, 80, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(80, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (8): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          80, 480, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          480, 480, kernel_size=(5, 5), stride=[1, 1], groups=480, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(480, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          480, 20, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          20, 480, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          480, 112, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(112, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (9): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          672, 672, kernel_size=(5, 5), stride=(1, 1), groups=672, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          672, 28, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          28, 672, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          672, 112, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(112, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (10): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          672, 672, kernel_size=(5, 5), stride=(1, 1), groups=672, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          672, 28, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          28, 672, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          672, 112, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(112, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (11): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          112, 672, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          672, 672, kernel_size=(5, 5), stride=[2, 2], groups=672, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 2, 1, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(672, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          672, 28, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          28, 672, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          672, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (12): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(5, 5), stride=(1, 1), groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (13): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(5, 5), stride=(1, 1), groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (14): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(5, 5), stride=(1, 1), groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(2, 2, 2, 2), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 192, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(192, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (15): MBConvBlock(\n",
       "        (_expand_conv): Conv2dStaticSamePadding(\n",
       "          192, 1152, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn0): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_depthwise_conv): Conv2dStaticSamePadding(\n",
       "          1152, 1152, kernel_size=(3, 3), stride=[1, 1], groups=1152, bias=False\n",
       "          (static_padding): ZeroPad2d(padding=(1, 1, 1, 1), value=0.0)\n",
       "        )\n",
       "        (_bn1): BatchNorm2d(1152, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "        (_se_reduce): Conv2dStaticSamePadding(\n",
       "          1152, 48, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_se_expand): Conv2dStaticSamePadding(\n",
       "          48, 1152, kernel_size=(1, 1), stride=(1, 1)\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_project_conv): Conv2dStaticSamePadding(\n",
       "          1152, 320, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (static_padding): Identity()\n",
       "        )\n",
       "        (_bn2): BatchNorm2d(320, eps=0.001, momentum=0.010000000000000009, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (decoder): UnetDecoder(\n",
       "    (layer1): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(432, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer2): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(296, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer3): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(152, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer4): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (layer5): DecoderBlock(\n",
       "      (attention1): Identity()\n",
       "      (attention2): Identity()\n",
       "      (block): Sequential(\n",
       "        (0): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "        (1): Conv2dReLU(\n",
       "          (block): Sequential(\n",
       "            (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (2): ReLU(inplace=True)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_conv): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (activation): Sigmoid()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unet = Unet('efficientnet-b0')\n",
    "unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 3, 3, 3])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "unet.encoder._conv_stem.weight.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "s = unet.encoder._conv_stem\n",
    "??s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 1, 3, 3])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "old_weights = unet.encoder._conv_stem.weight\n",
    "avg_weights = old_weights.mean(dim=1, keepdim=True)\n",
    "#avg_weights = avg_weights.unsqueeze(1)\n",
    "avg_weights.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
